package com.test.chatrobot3.test;

import java.util.*;

public class VocabularyBuilder {
    private Map<String, Integer> wordToIndex;
    private Map<Integer, String> indexToWord;
    private int nextIndex;

    public VocabularyBuilder() {
        wordToIndex = new HashMap<>();
        indexToWord = new HashMap<>();
        nextIndex = 0;
    }

    public void addWord(String word) {
        if (!wordToIndex.containsKey(word)) {
            wordToIndex.put(word, nextIndex);
            indexToWord.put(nextIndex, word);
            nextIndex++;
        }
    }

    public int getIndex(String word) {
        return wordToIndex.getOrDefault(word, -1);
    }

    public String getWord(int index) {
        return indexToWord.getOrDefault(index, "UNKNOWN");
    }

    public int size() {
        return wordToIndex.size();
    }

    public static void main(String[] args) {
        VocabularyBuilder vocab = new VocabularyBuilder();
        List<String> tokens = Arrays.asList("hello", "world", "hello", "test");
        for (String token : tokens) {
            vocab.addWord(token);
        }
        System.out.println(vocab.getIndex("hello"));  // 输出: 0
        System.out.println(vocab.getWord(1));         // 输出: world
    }
}